from typing import Dict, List, Type, Any
from datetime import datetime, timedelta
import os
import pandas as pd

from vnpy.trader.constant import Exchange, Interval, Direction, Offset
from vnpy.trader.object import BarData, OrderData, TradeData
from vnpy_ctastrategy import CtaTemplate
import vnpy_ctastrategy
from vnpy_ctastrategy.backtesting import BacktestingEngine

from portfolio_trader.datafeed.data_loader import PortfolioDataLoader
from portfolio_trader.signal.signal_result import SignalResult


class SignalGenerator:
    """量化信号生成器"""
    
    def __init__(self):
        """初始化"""
        self.data_loader = PortfolioDataLoader()
        self.strategy_classes: Dict[str, Type[CtaTemplate]] = {}
        self.strategy_settings: Dict[str, Dict] = {}
        self.signals: Dict[str, List[SignalResult]] = {}  # symbol: signal_list
        
        # 初始化策略类字典
        self._init_strategies()
        
    def _init_strategies(self) -> None:
        """初始化策略类"""
        # 导入策略模块
        from vnpy_ctastrategy.strategies.boll_channel_strategy import BollChannelStrategy
        from vnpy_ctastrategy.strategies.atr_rsi_strategy import AtrRsiStrategy
        from vnpy_ctastrategy.strategies.double_ma_strategy import DoubleMaStrategy
        from vnpy_ctastrategy.strategies.dual_thrust_strategy import DualThrustStrategy
        from vnpy_ctastrategy.strategies.king_keltner_strategy import KingKeltnerStrategy
        from vnpy_ctastrategy.strategies.multi_signal_strategy import MultiSignalStrategy
        from vnpy_ctastrategy.strategies.turtle_signal_strategy import TurtleSignalStrategy
        # 添加策略类到字典
        self.strategy_classes["BollChannel"] = BollChannelStrategy
        self.strategy_classes["AtrRsi"] = AtrRsiStrategy
        self.strategy_classes["DoubleMa"] = DoubleMaStrategy
        self.strategy_classes["DualThrust"] = DualThrustStrategy
        self.strategy_classes["KingKeltner"] = KingKeltnerStrategy
        self.strategy_classes["MultiSignal"] = MultiSignalStrategy
        self.strategy_classes["TurtleSignal"] = TurtleSignalStrategy
        # 设置默认参数
        # self.strategy_settings["BollChannel"] = {
        #     "boll_window": 18,
        #     "boll_dev": 3.4,
        #     "cci_window": 10,
        #     "atr_window": 30,
        #     "sl_multiplier": 5.2,
        #     "fixed_size": 100
        # }

    def _load_bar_data(self, symbol: str, start_date: datetime, end_date: datetime) -> List[BarData]:
        """加载K线数据"""
        return self.data_loader.load_daily_data(symbol, start_date, end_date)
    
    def _convert_to_strategy_bars(self, bars: List[BarData]) -> List[BarData]:
        """转换K线数据为策略可用格式"""
        # 按时间正序排列
        bars.sort(key=lambda x: x.datetime)
        return bars
    
    def generate_signal(self, symbol: str, strategy_name: str, days: int = 60) -> List[SignalResult]:
        """
        为指定股票生成量化信号
        
        参数:
            symbol (str): 股票代码
            strategy_name (str): 策略名称
            days (int, optional): 回溯天数. 默认为 60.
            
        返回:
            List[SignalResult]: 信号结果列表
        """
        # 检查策略是否存在
        if strategy_name not in self.strategy_classes:
            print(f"策略 {strategy_name} 不存在")
            return []
        
        # 获取日期范围
        end_date = datetime.now()
        start_date = end_date - timedelta(days=days)
        
        # 加载K线数据
        bars = self._load_bar_data(symbol, start_date, end_date)
        if not bars:
            print(f"未找到 {symbol} 的K线数据")
            return []
        
        # 转换K线数据
        strategy_bars = self._convert_to_strategy_bars(bars)
        
        # 创建回测引擎
        engine = BacktestingEngine()
        engine.set_parameters(
            vt_symbol=f"{symbol}.{bars[0].exchange.value}",
            interval=Interval.DAILY,
            start=start_date,
            end=end_date,
            rate=0.0003,  # 手续费率
            slippage=0.01,  # 滑点
            size=100,  # 合约乘数
            pricetick=0.01,  # 价格最小变动
            capital=1_000_000,  # 初始资金
        )
        
        # 添加策略
        strategy_class = self.strategy_classes[strategy_name]
        engine.add_strategy(strategy_class, {})
        
        # 加载数据
        engine.history_data.extend(strategy_bars)
        
        # 运行回测
        engine.run_backtesting()
        
        # 获取交易记录
        trades = engine.get_all_trades()
        
        # 生成信号结果
        signals = []
        
        # 检查策略是否有买入或卖出信号
        has_signal = False
        
        for trade in trades:
            signal = SignalResult(symbol, strategy_name)
            signal.datetime = trade.datetime
            
            # 根据交易方向和开平设置信号类型
            if trade.direction == Direction.LONG and trade.offset == Offset.OPEN:
                signal.set_buy_signal(trade.price, trade.volume)
            elif trade.direction == Direction.LONG and trade.offset == Offset.CLOSE:
                signal.set_sell_signal(trade.price, trade.volume)
            else:
                signal.set_hold_signal()
                
            signals.append(signal)
            has_signal = True
        
        # 如果没有生成任何信号，添加一个持仓信号
        if not has_signal and strategy_bars:
            signal = SignalResult(symbol, strategy_name)
            signal.datetime = strategy_bars[-1].datetime  # 使用最新K线的时间
            signal.set_hold_signal()
            signals.append(signal)
        
        # 保存信号
        self.signals[symbol] = signals
        
        return signals
    
    def generate_portfolio_signals(self, symbols: List[str], strategy_name: str, days: int = 365) -> Dict[str, List[SignalResult]]:
        """为投资组合生成交易信号
        
        参数:
            symbols (List[str]): 股票代码列表
            strategy_name (str): 策略名称
            days (int): 回溯天数
            
        返回:
            Dict[str, List[SignalResult]]: 股票代码到信号列表的映射
        """
        result = {}
        for symbol in symbols:
            signals = self.generate_signal(symbol, strategy_name, days)
            if signals:
                result[symbol] = signals
        return result
    
    def get_latest_signals(self, symbols: List[str], strategy_name: str) -> Dict[str, SignalResult]:
        """获取最新的交易信号
        
        参数:
            symbols (List[str]): 股票代码列表
            strategy_name (str): 策略名称
            
        返回:
            Dict[str, SignalResult]: 股票代码到最新信号的映射
        """
        # 先生成信号
        self.generate_portfolio_signals(symbols, strategy_name)
        
        # 获取每个股票的最新信号
        latest_signals = {}
        for symbol in symbols:
            if symbol in self.signals and self.signals[symbol]:
                # 按时间排序，获取最新的信号
                sorted_signals = sorted(self.signals[symbol], key=lambda x: x.datetime, reverse=True)
                latest_signals[symbol] = sorted_signals[0]
        
        return latest_signals